import numpy as np
from scipy import linalg, special
from traits.api import HasTraits, Array, CFloat, List, \
   Instance, on_trait_change, Property
from traitsui.api import Item, View, ListEditor, \
        HSplit, VSplit
from mayavi.core.ui.api import EngineView, MlabSceneModel, \
        SceneEditor
from traits.api import *
from traitsui.api import *
from mayavi.core.api import Engine
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
from numpy import sin, cos, mgrid, pi, sqrt

X, Y, Z = np.mgrid[-0.15:0.15:20j, -0.15:0.15:20j, -0.15:0.15:20j]
f = 1e4
X = np.round(X * f) / f
Y = np.round(Y * f) / f
Z = np.round(Z * f) / f

class Form(HasTraits):
    engine = Instance(Engine, ())
    scene = Instance(MlabSceneModel)
    def _scene_default(self):
        self.engine.start()
        scene = MlabSceneModel(engine=self.engine)
        return scene

    direction = Array(float, value=(0, 0, 1), cols=3,
                    shape=(3,), desc='directing vector of the loop',
                    enter_set=True, auto_set=False)
    radius    = CFloat(0.1, desc='radius of the loop',
                    enter_set=True, auto_set=False)
    position  = Array(float, value=(0, 0, 0), cols=3,
                    shape=(3,), desc='position of the center of the loop',
                    enter_set=True, auto_set=False)
    _plot      = None
    Bnorm   = Property(depends_on='direction,position,radius')
    view = View('position', 'direction', 'radius', '_')

    def base_vectors(self):
        n = self.direction / (self.direction**2).sum(axis=-1)
        if  np.abs(n[0])==1 :
            l = np.r_[n[2], 0, -n[0]]
        else:
            l = np.r_[0, n[2], -n[1]]
        l /= (l**2).sum(axis=-1)
        m = np.cross(n, l)
        return n, l, m

    @observe('Bnorm')
    def redraw(self,a):
        if hasattr(self, 'app') and self.app.scene._renderer is not None:
            self.display()
            self.app.visualize_field()

    def display(self,a):
        n, l, m = self.base_vectors()
        theta = np.linspace(0, 2*np.pi, 30)[..., np.newaxis]
        coil = self.radius*(np.sin(theta)*l + np.cos(theta)*m)
        coil += self.position
        coil_x, coil_y, coil_z = coil.T
        if self._plot is None:
            self._plot = self.app.scene.mlab.plot3d(coil_x, coil_y, coil_z,
                                    tube_radius=0.007, color=(0, 0, 1),
                                    name='Coil')
        else:
            self._plot.mlab_source.trait_set(x=coil_x, y=coil_y, z=coil_z)

    def _get_Bnorm(self,a):
        n, l, m = self.base_vectors()
        R       = self.radius
        r0      = self.position
        r       = np.c_[np.ravel(X), np.ravel(Y), np.ravel(Z)]
        trans = np.vstack((l, m, n))
        r -= r0
        r = np.dot(r, linalg.inv(trans) )
        x = r[:, 0]
        y = r[:, 1]
        z = r[:, 2]
        rho = np.sqrt(x**2 + y**2)
        theta = np.arctan2(x, y)
        E = special.ellipe((4 * R * rho)/( (R + rho)**2 + z**2))
        K = special.ellipk((4 * R * rho)/( (R + rho)**2 + z**2))
        Bz =  1/np.sqrt((R + rho)**2 + z**2) * (
                    K
                  + E * (R**2 - rho**2 - z**2)/((R - rho)**2 + z**2)
                )
        Brho = z/(rho*np.sqrt((R + rho)**2 + z**2)) * (
                -K
                + E * (R**2 + rho**2 + z**2)/((R - rho)**2 + z**2)
                )

        Brho[np.isnan(Brho)] = 0
        B = np.c_[np.cos(theta)*Brho, np.sin(theta)*Brho, Bz ]
        B = np.dot(B, trans)
        Bx, By, Bz = B.T
        Bx = np.reshape(Bx, X.shape)
        By = np.reshape(By, X.shape)
        Bz = np.reshape(Bz, X.shape)
        Bnorm = np.sqrt(Bx**2 + By**2 + Bz**2)
        Bmax = 10 * np.median(Bnorm)
        Bx[Bnorm > Bmax] = np.NAN
        By[Bnorm > Bmax] = np.NAN
        Bz[Bnorm > Bmax] = np.NAN
        Bnorm[Bnorm > Bmax] = np.NAN
        self.Bx = Bx
        self.By = By
        self.Bz = Bz
        return Bnorm

from PyQt5.QtWidgets import *
class FormWidget(QWidget):

    scene = Instance(MlabSceneModel, (), editor=SceneEditor())
    engine_view = Instance(EngineView)
    coils = List(Instance(Form, (), allow_none=False),
                        editor=ListEditor(style='custom'),
                        value=[ Form(position=(0, 0, -0.05), ),
                                 Form(position=(0, 0,  0.05), ), ])
    Bx    = Array(value=np.zeros_like(X))
    By    = Array(value=np.zeros_like(X))
    Bz    = Array(value=np.zeros_like(X))
    Bnorm = Array(value=np.zeros_like(X))
    vector_field = None

    def __init__(self, **traits):
        HasTraits.__init__(self, **traits)
        self.engine_view = EngineView(engine=self.scene.engine)

    @observe('scene.activated,coils')
    def init_view(self,a):
        if self.scene._renderer is not None:
            self.scene.scene_editor.background = (0, 0, 0)
            for coil in self.coils:
                coil.app = self
                coil.display()
            self.visualize_field()

    def visualize_field(self,a):
        self.Bx    = np.zeros_like(X)
        self.By    = np.zeros_like(X)
        self.Bz    = np.zeros_like(X)
        self.Bnorm = np.zeros_like(X)
        self.scene.scene.disable_render = True
        for coil in self.coils:
            self.Bx += coil.Bx
            self.By += coil.By
            self.Bz += coil.Bz
        self.Bnorm = np.sqrt(self.Bx**2 + self.By**2 + self.Bz**2)
        if self.vector_field is None:
            self.vector_field = self.scene.mlab.pipeline.vector_field(
                                    X, Y, Z, self.Bx, self.By, self.Bz,
                                    scalars=self.Bnorm,
                                    name='B field')
            vectors = self.scene.mlab.pipeline.vectors(self.vector_field,
                                    mode='arrow', resolution=10,
                                    mask_points=6, colormap='YlOrRd',
                                    scale_factor=2*np.abs(X[0,0,0]
                                                          -X[1,1,1]) )
            vectors.module_manager.vector_lut_manager.reverse_lut = True
            vectors.glyph.mask_points.random_mode = False
            self.scene.mlab.axes()
            self.scp = self.scene.mlab.pipeline.scalar_cut_plane(
                                                      self.vector_field,
                                                      colormap='hot')
        else:
            self.vector_field.mlab_source.trait_set(u=self.Bx, v=self.By, w=self.Bz,
                                              scalars=self.Bnorm)
        self.scene.scene.disable_render = False

    view = View(HSplit(
                    VSplit(Item(name='engine_view',
                                   style='custom',
                                   resizable=True),
                            Item('coils', springy=True),
                        show_labels=False),
                        'scene',
                        show_labels=False),
                    resizable=True,
                    title='Coils...',
                    height=0.8,
                    width=0.8,
                )

    def __init__(self):
        QWidget.__init__(self)
        layout = QHBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        self.form = Form()
        layout.addWidget(self.form.edit_traits(parent=self, kind='subpanel').control)
formWidget = FormWidget()
if __name__ == '__main__':
    import os
    import sys

    app = QApplication(sys.argv)
    formWidget.show()
    sys.exit(app.exec_())
#
# if __name__ == '__main__':
#     app = Application()
#     app.configure_traits()

